# -*- coding: utf-8 -*-
"""
基础文档分块器
提供通用的文档分块功能和UI界面
"""

import sys
import re
import docx2txt
import tkinter as tk
from tkinter import filedialog, messagebox, ttk
from tkinter.scrolledtext import ScrolledText
import threading
import time
import json
import os
from abc import ABC, abstractmethod
from docx import Document
from typing import List, Dict, Any, Optional


class ChunkingStrategy(ABC):
    """分块策略抽象基类"""
    
    @abstractmethod
    def chunk_text(self, text: str, **kwargs) -> List[str]:
        """分块文本的抽象方法
        
        Args:
            text: 待分块的文本
            **kwargs: 分块参数
            
        Returns:
            分块后的文本列表
        """
        pass


class SeparatorChunkingStrategy(ChunkingStrategy):
    """基于分隔符的分块策略"""
    
    def chunk_text(self, text: str, separator: str = "#", min_length: int = 100, **kwargs) -> List[str]:
        """基于分隔符进行文本分块
        
        Args:
            text: 待分块的文本
            separator: 分隔符
            min_length: 最小块长度
            **kwargs: 其他参数
            
        Returns:
            分块后的文本列表
        """
        chunks = []
        parts = text.split(separator)
        
        for part in parts:
            part = part.strip()
            if len(part) >= min_length:
                chunks.append(part)
        
        return chunks


class SmartChunkingStrategy(ChunkingStrategy):
    """智能分块策略"""
    
    def chunk_text(self, text: str, max_length: int = 1000, overlap: int = 100, **kwargs) -> List[str]:
        """智能分块，考虑段落和句子边界
        
        Args:
            text: 待分块的文本
            max_length: 最大块长度
            overlap: 重叠长度
            **kwargs: 其他参数
            
        Returns:
            分块后的文本列表
        """
        chunks = []
        paragraphs = text.split('\n\n')
        current_chunk = ""
        
        for paragraph in paragraphs:
            if len(current_chunk) + len(paragraph) <= max_length:
                current_chunk += paragraph + "\n\n"
            else:
                if current_chunk:
                    chunks.append(current_chunk.strip())
                current_chunk = paragraph + "\n\n"
        
        if current_chunk:
            chunks.append(current_chunk.strip())
        
        return chunks


class BaseChunker:
    """基础文档分块器类"""
    
    def __init__(self, master: tk.Tk, title: str = "文档分块工具"):
        """初始化基础分块器
        
        Args:
            master: Tkinter主窗口
            title: 窗口标题
        """
        self.master = master
        self.master.title(title)
        self.master.geometry("800x600")
        self.master.resizable(True, True)
        
        # 分块策略
        self.strategies = {
            "按分隔符": SeparatorChunkingStrategy(),
            "智能分块": SmartChunkingStrategy()
        }
        
        # 文件列表
        self.file_list = []
        self.processing_thread = None
        self.stop_processing = False
        
        # 初始化UI
        self._init_ui()
        
    def _init_ui(self):
        """初始化用户界面"""
        # 设置主窗口样式
        self.style = ttk.Style()
        self.style.configure('TButton', font=('微软雅黑', 10))
        self.style.configure('TLabel', font=('微软雅黑', 10))
        
        # 创建主框架
        self.main_frame = ttk.Frame(self.master, padding="10")
        self.main_frame.pack(fill=tk.BOTH, expand=True)
        
        # 创建顶部操作区域
        self.top_frame = ttk.Frame(self.main_frame)
        self.top_frame.pack(fill=tk.X, padx=5, pady=5)
        
        self._create_file_selection_frame()
        self._create_output_frame()
        self._create_parameter_frame()
        self._create_control_frame()
        self._create_status_frame()
        self._create_log_frame()
        
    def _create_file_selection_frame(self):
        """创建文件选择框架"""
        self.file_frame = ttk.LabelFrame(self.top_frame, text="文件选择", padding="5")
        self.file_frame.pack(fill=tk.X, padx=5, pady=5)
        
        # 文件路径输入框
        self.path_var = tk.StringVar()
        self.path_entry = ttk.Entry(self.file_frame, textvariable=self.path_var, width=50)
        self.path_entry.pack(side=tk.LEFT, fill=tk.X, expand=True, padx=5, pady=5)
        
        # 选择文件按钮
        self.select_file_btn = ttk.Button(self.file_frame, text="选择文件", command=self.select_file)
        self.select_file_btn.pack(side=tk.LEFT, padx=5, pady=5)
        
        # 批量选择文件按钮
        self.select_multiple_files_btn = ttk.Button(self.file_frame, text="批量选择", command=self.select_multiple_files)
        self.select_multiple_files_btn.pack(side=tk.LEFT, padx=5, pady=5)
        
        # 选择文件夹按钮
        self.select_folder_btn = ttk.Button(self.file_frame, text="选择文件夹", command=self.select_folder)
        self.select_folder_btn.pack(side=tk.LEFT, padx=5, pady=5)
        
    def _create_output_frame(self):
        """创建输出目录框架"""
        self.output_frame = ttk.LabelFrame(self.top_frame, text="输出目录", padding="5")
        self.output_frame.pack(fill=tk.X, padx=5, pady=5)
        
        # 输出目录输入框
        self.output_var = tk.StringVar()
        self.output_entry = ttk.Entry(self.output_frame, textvariable=self.output_var, width=50)
        self.output_entry.pack(side=tk.LEFT, fill=tk.X, expand=True, padx=5, pady=5)
        
        # 选择输出目录按钮
        self.select_output_btn = ttk.Button(self.output_frame, text="选择输出目录", command=self.select_output_dir)
        self.select_output_btn.pack(side=tk.LEFT, padx=5, pady=5)
        
    def _create_parameter_frame(self):
        """创建分块参数框架"""
        self.param_frame = ttk.LabelFrame(self.top_frame, text="分块参数", padding="5")
        self.param_frame.pack(fill=tk.X, padx=5, pady=5)
        
        # 分块模式
        ttk.Label(self.param_frame, text="分块模式:").pack(side=tk.LEFT, padx=5, pady=5)
        self.chunk_mode_var = tk.StringVar(value="按分隔符")
        self.chunk_mode_combo = ttk.Combobox(self.param_frame, textvariable=self.chunk_mode_var, width=15)
        self.chunk_mode_combo['values'] = tuple(self.strategies.keys())
        self.chunk_mode_combo.pack(side=tk.LEFT, padx=5, pady=5)
        
        # 分隔符设置
        ttk.Label(self.param_frame, text="分隔符:").pack(side=tk.LEFT, padx=5, pady=5)
        self.separator_var = tk.StringVar(value="#")
        self.separator_entry = ttk.Entry(self.param_frame, textvariable=self.separator_var, width=5)
        self.separator_entry.pack(side=tk.LEFT, padx=5, pady=5)
        
        # 最小内容长度
        ttk.Label(self.param_frame, text="最小内容长度:").pack(side=tk.LEFT, padx=5, pady=5)
        self.min_content_length_var = tk.IntVar(value=100)
        self.min_content_length_spinbox = ttk.Spinbox(self.param_frame, from_=50, to=500, increment=50, textvariable=self.min_content_length_var, width=5)
        self.min_content_length_spinbox.pack(side=tk.LEFT, padx=5, pady=5)
        
        # 超时设置
        ttk.Label(self.param_frame, text="超时(秒):").pack(side=tk.LEFT, padx=5, pady=5)
        self.timeout_var = tk.IntVar(value=30)
        self.timeout_spinbox = ttk.Spinbox(self.param_frame, from_=10, to=300, increment=10, textvariable=self.timeout_var, width=5)
        self.timeout_spinbox.pack(side=tk.LEFT, padx=5, pady=5)
        
    def _create_control_frame(self):
        """创建控制按钮框架"""
        self.control_frame = ttk.Frame(self.top_frame)
        self.control_frame.pack(fill=tk.X, padx=5, pady=5)
        
        # 开始分块按钮
        self.start_btn = ttk.Button(self.control_frame, text="开始分块", command=self.start_chunking)
        self.start_btn.pack(side=tk.LEFT, padx=5, pady=5)
        
        # 停止分块按钮
        self.stop_btn = ttk.Button(self.control_frame, text="停止分块", command=self.stop_chunking, state=tk.DISABLED)
        self.stop_btn.pack(side=tk.LEFT, padx=5, pady=5)
        
        # 清空日志按钮
        self.clear_log_btn = ttk.Button(self.control_frame, text="清空日志", command=self.clear_log)
        self.clear_log_btn.pack(side=tk.LEFT, padx=5, pady=5)
        
    def _create_status_frame(self):
        """创建状态显示框架"""
        self.status_frame = ttk.Frame(self.main_frame)
        self.status_frame.pack(fill=tk.X, padx=5, pady=5)
        
        # 进度条
        self.progress_var = tk.DoubleVar()
        self.progress_bar = ttk.Progressbar(self.status_frame, variable=self.progress_var, maximum=100)
        self.progress_bar.pack(fill=tk.X, padx=5, pady=5)
        
        # 状态标签
        self.status_var = tk.StringVar(value="就绪")
        self.status_label = ttk.Label(self.status_frame, textvariable=self.status_var)
        self.status_label.pack(side=tk.LEFT, padx=5, pady=5)
        
        # 文件计数标签
        self.file_count_var = tk.StringVar(value="文件: 0/0")
        self.file_count_label = ttk.Label(self.status_frame, textvariable=self.file_count_var)
        self.file_count_label.pack(side=tk.RIGHT, padx=5, pady=5)
        
    def _create_log_frame(self):
        """创建日志显示框架"""
        self.log_frame = ttk.LabelFrame(self.main_frame, text="处理日志", padding="5")
        self.log_frame.pack(fill=tk.BOTH, expand=True, padx=5, pady=5)
        
        # 日志文本区域
        self.log_text = ScrolledText(self.log_frame, height=15, wrap=tk.WORD)
        self.log_text.pack(fill=tk.BOTH, expand=True)
        
    def select_file(self):
        """选择单个文件"""
        file_path = filedialog.askopenfilename(
            title="选择文档文件",
            filetypes=[("Word文档", "*.docx"), ("文本文件", "*.txt"), ("所有文件", "*.*")]
        )
        if file_path:
            self.path_var.set(file_path)
            self.file_list = [file_path]
            self.log(f"已选择文件: {file_path}")
            
    def select_multiple_files(self):
        """选择多个文件"""
        file_paths = filedialog.askopenfilenames(
            title="选择多个文档文件",
            filetypes=[("Word文档", "*.docx"), ("文本文件", "*.txt"), ("所有文件", "*.*")]
        )
        if file_paths:
            self.file_list = list(file_paths)
            self.path_var.set(f"已选择 {len(file_paths)} 个文件")
            self.log(f"已选择 {len(file_paths)} 个文件")
            
    def select_folder(self):
        """选择文件夹"""
        folder_path = filedialog.askdirectory(title="选择包含文档的文件夹")
        if folder_path:
            # 扫描文件夹中的文档文件
            self.file_list = []
            for root, dirs, files in os.walk(folder_path):
                for file in files:
                    if file.lower().endswith(('.docx', '.txt')):
                        self.file_list.append(os.path.join(root, file))
            
            if self.file_list:
                self.path_var.set(f"文件夹: {folder_path} ({len(self.file_list)} 个文件)")
                self.log(f"从文件夹 {folder_path} 中找到 {len(self.file_list)} 个文档文件")
            else:
                messagebox.showwarning("警告", "所选文件夹中没有找到支持的文档文件")
                
    def select_output_dir(self):
        """选择输出目录"""
        output_dir = filedialog.askdirectory(title="选择输出目录")
        if output_dir:
            self.output_var.set(output_dir)
            self.log(f"输出目录设置为: {output_dir}")
            
    def start_chunking(self):
        """开始分块处理"""
        if not self.file_list:
            messagebox.showerror("错误", "请先选择要处理的文件")
            return
            
        if not self.output_var.get():
            messagebox.showerror("错误", "请选择输出目录")
            return
            
        # 禁用开始按钮，启用停止按钮
        self.start_btn.config(state=tk.DISABLED)
        self.stop_btn.config(state=tk.NORMAL)
        self.stop_processing = False
        
        # 启动处理线程
        self.processing_thread = threading.Thread(target=self._process_files)
        self.processing_thread.daemon = True
        self.processing_thread.start()
        
    def stop_chunking(self):
        """停止分块处理"""
        self.stop_processing = True
        self.log("正在停止处理...")
        
    def clear_log(self):
        """清空日志"""
        self.log_text.delete(1.0, tk.END)
        
    def log(self, message: str):
        """记录日志信息
        
        Args:
            message: 日志消息
        """
        timestamp = time.strftime("%H:%M:%S")
        log_message = f"[{timestamp}] {message}\n"
        
        # 在主线程中更新UI
        self.master.after(0, lambda: self._update_log(log_message))
        
    def _update_log(self, message: str):
        """更新日志显示
        
        Args:
            message: 日志消息
        """
        self.log_text.insert(tk.END, message)
        self.log_text.see(tk.END)
        
    def _process_files(self):
        """处理文件的主要逻辑"""
        try:
            total_files = len(self.file_list)
            processed_files = 0
            
            for i, file_path in enumerate(self.file_list):
                if self.stop_processing:
                    break
                    
                # 更新状态
                self.master.after(0, lambda: self.status_var.set(f"正在处理: {os.path.basename(file_path)}"))
                self.master.after(0, lambda: self.file_count_var.set(f"文件: {i+1}/{total_files}"))
                
                try:
                    self._process_single_file(file_path)
                    processed_files += 1
                    
                    # 更新进度
                    progress = (i + 1) / total_files * 100
                    self.master.after(0, lambda p=progress: self.progress_var.set(p))
                    
                except Exception as e:
                    self.log(f"处理文件 {file_path} 时出错: {str(e)}")
                    
            # 处理完成
            if not self.stop_processing:
                self.log(f"处理完成！共处理 {processed_files}/{total_files} 个文件")
                self.master.after(0, lambda: self.status_var.set("处理完成"))
            else:
                self.log("处理已停止")
                self.master.after(0, lambda: self.status_var.set("已停止"))
                
        except Exception as e:
            self.log(f"处理过程中发生错误: {str(e)}")
            
        finally:
            # 恢复按钮状态
            self.master.after(0, lambda: self.start_btn.config(state=tk.NORMAL))
            self.master.after(0, lambda: self.stop_btn.config(state=tk.DISABLED))
            
    def _process_single_file(self, file_path: str):
        """处理单个文件
        
        Args:
            file_path: 文件路径
        """
        self.log(f"开始处理文件: {os.path.basename(file_path)}")
        
        # 读取文件内容
        if file_path.lower().endswith('.docx'):
            text = docx2txt.process(file_path)
        else:
            with open(file_path, 'r', encoding='utf-8') as f:
                text = f.read()
                
        if not text.strip():
            self.log(f"文件 {file_path} 内容为空，跳过处理")
            return
            
        # 获取分块策略和参数
        strategy_name = self.chunk_mode_var.get()
        strategy = self.strategies[strategy_name]
        
        # 准备分块参数
        chunk_params = {
            'separator': self.separator_var.get(),
            'min_length': self.min_content_length_var.get(),
            'max_length': 1000,
            'overlap': 100
        }
        
        # 执行分块
        chunks = strategy.chunk_text(text, **chunk_params)
        
        if not chunks:
            self.log(f"文件 {file_path} 分块后无有效内容")
            return
            
        # 保存分块结果
        self._save_chunks(file_path, chunks)
        
        self.log(f"文件 {os.path.basename(file_path)} 处理完成，生成 {len(chunks)} 个分块")
        
    def _save_chunks(self, original_file_path: str, chunks: List[str]):
        """保存分块结果
        
        Args:
            original_file_path: 原始文件路径
            chunks: 分块列表
        """
        output_dir = self.output_var.get()
        base_name = os.path.splitext(os.path.basename(original_file_path))[0]
        
        # 创建输出文件路径
        output_file = os.path.join(output_dir, f"{base_name}_processed.txt")
        
        # 确保输出目录存在
        os.makedirs(output_dir, exist_ok=True)
        
        # 保存分块内容
        with open(output_file, 'w', encoding='utf-8') as f:
            for i, chunk in enumerate(chunks, 1):
                f.write(f"# 分块 {i}\n\n")
                f.write(chunk)
                f.write("\n\n" + "="*50 + "\n\n")
                
        self.log(f"分块结果已保存到: {output_file}")